

import sys
import numpy as np
import scipy.signal as signal
import soundfile as sf
from frequency_domain_adaptive_filters.pfdkf import PFDKF


class ring_buffer:
    def __init__(self, size, delay):
        self.buffer = np.zeros(size)
        self.size = size
        self.head = 0
        self.tail = (size + delay) % size

    def get(self, length):
        if (self.tail + length) <= self.size:
            block = self.buffer[self.tail:self.tail+length]
        else:
            block = np.concatenate((self.buffer[self.tail:], self.buffer[:self.tail+length-self.size]))
        self.tail = (self.tail + length) % self.size
        return block

    def put(self, data):
        length = len(data)
        if (self.head + length) <= self.size:
            self.buffer[self.head:self.head+length] = data
        else:
            self.buffer[self.head:] = data[:self.size - self.head]
            self.buffer[:self.head+length-self.size] = data[self.size - self.head:]
        self.head = (self.head + length) % self.size



def RMS(x):
    return 10 * np.log10(np.average(x**2) + 1e-6)


IR = [
    7.3281956e-07,  -2.3036519e-06, -1.0323231e-06, -1.0247837e-06,
    2.212295e-06,   -7.415722e-07,  1.1829076e-06,  -4.745906e-07,
    2.9874734e-06,  -4.281139e-06,  8.586982e-08,   -6.501403e-07,
    2.2047632e-06,  -4.2968436e-06, 1.0665299e-06,  1.2596523e-06,
    5.412453e-06,   -4.689883e-06,  -3.5659642e-07, 3.765102e-06,
    1.7805148e-06,  -4.4947487e-06, -1.3232268e-06, 3.1110858e-06,
    2.5975182e-06,  -1.078435e-06,  5.435293e-07,   9.365622e-06,
    -2.0501457e-06, -6.554812e-06,  -2.0207767e-06, 3.3305216e-06,
    -1.1186148e-06, -4.1729963e-06, 1.6854503e-06,  8.955039e-06,
    -5.500982e-07,  -4.441023e-06,  6.69175e-06,    8.4743015e-06,
    -4.380904e-06,  -7.309603e-06,  -5.537905e-06,  6.038217e-06,
    -2.301219e-06,  7.1914546e-06,  1.7591914e-06,  6.9735893e-06,
    -2.2435363e-05, -8.356573e-06,  -1.4786573e-05, 5.6685712e-06,
    -1.7222818e-06, 3.554604e-05,   -1.7430219e-05, -1.4412982e-05,
    -7.088181e-05,  5.5517354e-05,  0.00011475379,  -4.3006567e-06,
    -0.00013208372, -0.00013692162, 0.00012412298,  0.0001520886,
    0.00011527779,  -0.00019172326, -3.839868e-05,  3.422383e-05,
    -1.3705036e-05, 9.911429e-06,   -3.8974426e-06, 9.182771e-06,
    -7.5424255e-06, 4.094613e-06,   -6.7274595e-06, 8.0840555e-06,
    -4.9138384e-06, 4.373888e-06,   -4.453825e-06,  5.04362e-06,
    -5.4085776e-06, 4.6573846e-06,  -3.994657e-06,  4.245168e-06,
    -4.8176244e-06, 2.9306561e-06,  -3.073625e-06,  3.4753916e-06,
    -3.0488482e-06, 2.7297738e-06,  -2.8932604e-06, 2.7092522e-06,
    -1.9935642e-06, 1.87816e-06,    -1.9920349e-06, 1.4089518e-06
]


N = 6
M = 32
A = 0.9
P_initial = 1e+2
partial_constrain = True

ft = PFDKF(N, M, A, P_initial, partial_constrain)


d, fs = sf.read('in.wav')

num_block = len(d) // M
e = np.zeros(num_block * M)
x = np.zeros(num_block * M)
feedback = np.zeros(num_block * M)
P = np.zeros(num_block * M)
dW = np.zeros(num_block * M)


out_buf = ring_buffer(fs, -2 * fs // 1000)
fb_buf = ring_buffer(fs, -2 * fs // 1000)


gain = 1 if len(sys.argv) < 2 else float(sys.argv[1])
scale = np.power(10, gain / 20)

fb_scale = np.power(10, 50 / 20)

IR = np.array(IR[::-1]) * fb_scale

np.save('ir.npy', IR)

zi = np.zeros(len(IR) - 1)

w = np.concatenate((IR, np.zeros(N * M - len(IR))))
W = np.fft.rfft(w.reshape((-1, M)), 2*M, axis=1)

d_rms = []
e_rms = []

h = np.zeros((num_block, N*M))


bypass =  len(sys.argv) >= 3


for n in range(num_block):
    x_n = out_buf.get(M)
    fb_n = fb_buf.get(M)
    d_n = d[n*M:(n+1)*M] + fb_n
    # d_n = d[n*M:(n+1)*M]

    d_rms.append(RMS(d_n))

    if bypass:
        e_n = d_n
    else:
        e_n = ft.filt(x_n, d_n)
        ft.update(e_n)

    e_rms.append(RMS(e_n))

    e_n = np.clip(e_n * scale, -1, 1)

    out_buf.put(e_n)

    fb, zi = signal.lfilter(IR, 1, e_n, zi=zi)
    fb_buf.put(fb)


    e[n*M:(n+1)*M] = e_n
    x[n*M:(n+1)*M] = x_n
    feedback[n*M:(n+1)*M] = fb_n
    P[n*M:(n+1)*M] = np.sum(ft.P)
    dW[n*M:(n+1)*M] = np.sum(np.abs(W - ft.H) ** 2)

    # if n == (5 * fs // M):
    #     scale = np.power(10, 40 / 20)

    h[n, :] = ft.h


np.save('h.npy', h)

mix = np.clip(np.array((e, P, dW)).T, -1, 1)


sf.write('feedback.wav', feedback, fs, 'PCM_16')
sf.write('out.wav', x, fs, 'PCM_16')
sf.write('mix.wav', mix, fs, 'PCM_16')

sys.exit(0)

import matplotlib.pyplot as plt
erle = np.array(d_rms) - np.array(e_rms)
plt.plot(erle)
plt.show()